from ultralytics import YOLO

if __name__ == '__main__':
    # 第一阶段配置
    stage1_config = {
        "data": "ultralytics/cfg/datasets/helmet3.yaml",
        "epochs": 500,
        "imgsz": 640,
        "cache": True,
        "batch": 32,
        "optimizer": "AdamW",
        "lr0": 0.01,
        "patience": 100,
        "device": 0,
        "project": "runs/detect",
        "workers": 0,
        "cos_lr": True,  # 启用余弦退火学习率
        "weight_decay": 0.0005,
        "warmup_epochs": 5,
        "warmup_momentum": 0.8,
        "dropout": 0.1,
        "label_smoothing": 0.05,
        "erasing": 0.4
    }

    # 执行第一阶段训练
    model = YOLO("weights/yolov8s.pt")
    model.train(**stage1_config)